#!/usr/bin/env python
# vim: set fileencoding=utf-8 :

import argparse

## XML FILE
import urllib2
from xml.dom.minidom import parseString
import re
## XML FILE

## Wave File
from scipy.io import wavfile

import codecs

def getSegments(nFile):
	file = open(nFile)
	xml = file.read()
	file.close()
	segments = []
	startTime = 'startTime=\"(.*)\"' # Fazer essa REGEX com um OR e pegar no
	stopTime = 'endTime=\"(.*)\" sp' # group
	speaker = 'speaker="(.*)" '
	Phrase = '>(.*)<'
	
	dom = parseString(xml).getElementsByTagName('UNIT')
	
	for i in range(len(dom)):
		srT = re.search(startTime,dom[i].toxml())
		stT =  re.search(stopTime,dom[i].toxml())
		speak = re.search(speaker,dom[i].toxml())
		ph =  re.search(Phrase,dom[i].toxml())
		segments.append( [str(speak.group(1)),
		float(srT.group(1)),float(stT.group(1)), ph.group(1)] )
	return segments

def creatSegment(input_wav, segments, output_wav):
	fs, data = wavfile.read(input_wav)
	Speaker = {}
	for i in range(len(segments)):
		if segments[i][0] not in Speaker:
			Speaker[segments[i][0]] = 1
		else:
			Speaker[segments[i][0]] += 1
	cSpeaker = Speaker.copy()
	for i in range(len(segments)):
		cSpeaker[segments[i][0]] -= 1
		wavfile.write(output_wav+'/%s-%03d-%d.wav' % (input_wav.replace('.wav',''), i+1, len(segments)), fs, data[segments[i][1]*fs:segments[i][2]*fs])
		wavfile.write(output_wav+'/%s-%03d-%d.wav' % (segments[i][0],
		Speaker[segments[i][0]] - cSpeaker[segments[i][0]],
		Speaker[segments[i][0]]), fs, data[segments[i][1]*fs:segments[i][2]*fs])

def transcription(segments, input_xml, output_wav):
	Speaker = {}
	allTransc = ''
	for i in range(len(segments)):
		if segments[i][0] not in Speaker:
			Speaker[segments[i][0]] = ''
	cSpeaker = Speaker.copy()
	for i in segments:
		cSpeaker[i[0]] += i[3]+'\n'
		allTransc += i[3]+'\n'
	for i in cSpeaker.items():
		out_file = codecs.open(output_wav+'/'+i[0]+'.txt', 'w','utf-8')
		out_file.write(unicode(i[1]))
		out_file.close()
	out_file = codecs.open(output_wav+'/'+input_xml.replace('.xml','')+'.txt',
	'w','utf-8')
	out_file.write(unicode(allTransc))
	out_file.close()

if __name__=='__main__':
	nfile = 'bpubmn01.xml'
	parser = argparse.ArgumentParser(description='Segmentation tool from \
	UFMG\'s corpora.')
	parser.add_argument('-i','--input-xml', required=True, help='read xml')
	parser.add_argument('-w','--input-wav', required=True, help='read wave')
	parser.add_argument('-o','--output-wav', required=False, help='output \
	wave', default = '.')
	args = parser.parse_args()
	seg = getSegments(args.input_xml)
	creatSegment(args.input_wav,seg, args.output_wav)
	transcription(seg, args.input_xml, args.output_wav)
